﻿using UnityEngine;
using System.Collections;
using System.Collections.Generic;
using MemoryProfilerWindow;
using System.Linq;
using UnityEditor;
using UnityEditor.MemoryProfiler;
using System.IO;

public class GroupMemSnapshot
{
    string filePath = "";
    CrawledMemorySnapshot cms = null;
    Dictionary<string, List<ThingInMemory>> nativeDic = new Dictionary<string, List<ThingInMemory>>();
    Dictionary<string, List<ThingInMemory>> manageDic = new Dictionary<string, List<ThingInMemory>>();

    List<KeyValuePair<string, long>> nativeSizeList = new List<KeyValuePair<string, long>>();
    List<KeyValuePair<string, long>> manageSizeList = new List<KeyValuePair<string, long>>();

    public GroupMemSnapshot(string file)
    {
        try
        {
            PackedMemorySnapshot pms = PackedMemorySnapshotUtility.LoadFromFile(file);
            filePath = file;
            EditorUtility.DisplayProgressBar("load snapshot", "Crawling Snapshot " + file + "...", 0.33f);
            PackedCrawlerData pcd = new Crawler().Crawl(pms);
            EditorUtility.DisplayProgressBar("load snapshot", "Unpacking Snapshot " + file + "...", 0.67f);
            cms = CrawlDataUnpacker.Unpack(pcd);
            EditorUtility.ClearProgressBar();

            foreach (ThingInMemory nativeObj in cms.nativeObjects)
            {
                string groupName = GetGroupName(nativeObj);
                groupName = string.IsNullOrEmpty(groupName) ? "UnknownNative" : groupName;

                if (!nativeDic.ContainsKey(groupName))
                {
                    nativeDic.Add(groupName, new List<ThingInMemory>());
                }

                nativeDic[groupName].Add(nativeObj);
            }

            foreach (ThingInMemory manageObj in cms.managedObjects)
            {
                string groupName = GetGroupName(manageObj);
                groupName = string.IsNullOrEmpty(groupName) ? "UnknownManage" : groupName;

                if (!manageDic.ContainsKey(groupName))
                {
                    manageDic.Add(groupName, new List<ThingInMemory>());
                }

                manageDic[groupName].Add(manageObj);
            }

            foreach (var kvp in nativeDic)
            {
                long size = kvp.Value.ToArray().Sum(o => o.size);
                nativeSizeList.Add(new KeyValuePair<string, long>(kvp.Key, size));
            }
            nativeSizeList.Sort((x, y) => -x.Value.CompareTo(y.Value));

            foreach (var kvp in manageDic)
            {
                long size = kvp.Value.ToArray().Sum(o => o.size);
                manageSizeList.Add(new KeyValuePair<string, long>(kvp.Key, size));
            }
            manageSizeList.Sort((x, y) => -x.Value.CompareTo(y.Value));
        }
        catch (System.Exception e)
        {
            Debug.Log("load snapshot err: "+e.ToString());
            EditorUtility.ClearProgressBar();
        }

		/*
        Debug.LogError("snapshot file: "+file);
        Debug.LogWarning("native count: "+nativeSizeList.Count);
        foreach (var item in nativeSizeList)
        {
            Debug.Log(item.Key+": "+formatBytes(item.Value));
        }
        Debug.LogWarning("manage count: "+manageSizeList.Count);
        foreach (var item in manageSizeList)
        {
            Debug.Log(item.Key+": "+formatBytes(item.Value));
        }
		*/
    }

	void diffClassSize(List<KeyValuePair<string, long>> sizeA
						, List<KeyValuePair<string, long>> sizeB
						, long classSizeFilter 
						, ref List<KeyValuePair<string, long>> diffSizeResult)
	{
		foreach (var item in sizeA)
		{
			long size1 = 0;
			if (sizeB.Exists(p=>p.Key == item.Key))
			{
				size1 = sizeB.Find(p => p.Key == item.Key).Value;
			}

			long diffSize = item.Value - size1;

			diffSizeResult.Add(new KeyValuePair<string, long>(item.Key, diffSize));
		}
		foreach (var item in sizeB)
		{
			if (diffSizeResult.Exists(p => p.Key == item.Key))
				continue;

			if (item.Value < classSizeFilter)
				continue;

			diffSizeResult.Add(new KeyValuePair<string, long>(item.Key, -item.Value));
		}
		diffSizeResult.Sort((x, y) => -x.Value.CompareTo(y.Value));
	}

	string diffClassAndItemSizeStr(List<KeyValuePair<string, long>> diffSizeList
									, long classSizeFilter
									, long itemSizeFilter
									, Dictionary<string, List<ThingInMemory>> classSizeDicA
									, Dictionary<string, List<ThingInMemory>> classSizeDicB
                                    , string classFilter)
	{
		string str = "";
		foreach (var item in diffSizeList)
		{
            if (!item.Key.Contains(classFilter))
                continue;

			int count1 = 0;
			int count2 = 0;
			List<ThingInMemory> itemListA = new List<ThingInMemory> ();
			List<ThingInMemory> itemListB = new List<ThingInMemory> ();

			if (classSizeDicA.ContainsKey (item.Key)) 
			{
				count1 = classSizeDicA [item.Key].Count;
				itemListA = new List<ThingInMemory> (classSizeDicA[item.Key]);
				itemListA.Sort ((x,y)=>-x.size.CompareTo(y.size));
			}
			if (classSizeDicB.ContainsKey (item.Key)) 
			{
				count2 = classSizeDicB [item.Key].Count;
				itemListB = new List<ThingInMemory> (classSizeDicB[item.Key]);
				itemListB.Sort ((x,y)=>-x.size.CompareTo(y.size));
			}

			long size = item.Value > 0 ? item.Value : item.Value * -1;
			if (size < classSizeFilter || 0 == item.Value)
				continue;
			
			str += item.Key + ": " +(count1-count2)+", "+ formatBytes(item.Value) + "\r\n";

			for (int i=0; i<itemListA.Count;)
			{
				var subitem = itemListA[i];
				itemListA.Remove (subitem);

				if (itemListB.Exists (p => p.caption == subitem.caption && p.size == subitem.size)) 
				{
					var subitemB = itemListB.Find (p => p.caption == subitem.caption && p.size == subitem.size);
					itemListB.Remove (subitemB);

					continue;
				}

				if (subitem.size < itemSizeFilter)
					continue;

				str += "\t[add] "+subitem.caption+": "+formatBytes(subitem.size)+"\r\n";
			}

			for (int i=0; i<itemListB.Count;)
			{
				var subitemB = itemListB[i];
				itemListB.Remove (subitemB);

				if (subitemB.size < itemSizeFilter)
					continue;
				
				str += "\t[rem] "+subitemB.caption +": "+formatBytes(subitemB.size) + "\r\n";
			}
		}
		return str;
	}

	string diffClassSizeStr(List<KeyValuePair<string, long>> diffSizeList
									, long classSizeFilter
									, Dictionary<string, List<ThingInMemory>> classSizeDicA
									, Dictionary<string, List<ThingInMemory>> classSizeDicB)
	{
		string str = "";
		foreach (var item in diffSizeList)
		{
			int count1 = 0;
			int count2 = 0;
			if (classSizeDicA.ContainsKey (item.Key)) 
			{
				count1 = classSizeDicA [item.Key].Count;
			}
			if (classSizeDicB.ContainsKey (item.Key)) 
			{
				count2 = classSizeDicB [item.Key].Count;
			}

			long size = item.Value > 0 ? item.Value : item.Value * -1;
			if (size < classSizeFilter || 0 == item.Value)
				continue;

			str += item.Key + ": " +(count1-count2)+", "+ formatBytes(item.Value) + "\r\n";
		}
		return str;
	}

	public void Diff(GroupMemSnapshot gms, long classSizeFilter, long itemSizeFilter, string classFilter, bool bNative, bool bManage)
	{
		if (0 > classSizeFilter)
			classSizeFilter = 0;
		if (0 > itemSizeFilter)
			itemSizeFilter = 0;

		List<KeyValuePair<string, long>> diffNativeSizeList = new List<KeyValuePair<string, long>>();
		List<KeyValuePair<string, long>> diffManageSizeList = new List<KeyValuePair<string, long>>();

		long diffTotalSize = cms.totalSize - gms.cms.totalSize;
		long diffNativeSize = cms.totalNativeObjSize - gms.cms.totalNativeObjSize;
		long diffManageSize = cms.totalManagedObjSize - gms.cms.totalManagedObjSize;

		diffClassSize (nativeSizeList, gms.nativeSizeList, classSizeFilter, ref diffNativeSizeList);
		diffClassSize (manageSizeList, gms.manageSizeList, classSizeFilter, ref diffManageSizeList);

        string str = "+++++++++++++++++++ outline diff +++++++++++++++\r\n";
        str += Path.GetFileNameWithoutExtension(filePath) + " - " + Path.GetFileNameWithoutExtension(gms.filePath) + ": \r\n";
        str += "total diff size: " +formatBytes(diffTotalSize)+"\r\n";
		str += "total diff native size: " + formatBytes(diffNativeSize) + "\r\n";
		str += "total diff manage size: " + formatBytes(diffManageSize) + "\r\n";

		str += "\r\n+++++++++++++++++++ outline snapshot +++++++++++++++\r\n";
		str += Path.GetFileNameWithoutExtension(filePath)+": "+outlineInfo()+"\r\n";
		str += Path.GetFileNameWithoutExtension(gms.filePath)+": "+gms.outlineInfo()+"\r\n";

        if (bNative)
        {
            str += "\r\n+++++++++++++++++++++ native ++++++++++++++++\r\n";
            str += diffClassAndItemSizeStr(diffNativeSizeList, classSizeFilter, itemSizeFilter, nativeDic, gms.nativeDic, classFilter);
        }

        if (bManage)
        {
            str += "\r\n+++++++++++++++++++++ manage ++++++++++++++++\r\n";
            str += diffClassAndItemSizeStr(diffManageSizeList, classSizeFilter, itemSizeFilter, manageDic, gms.manageDic, classFilter);
        }

		string outputFile = Path.GetFileNameWithoutExtension(filePath) + "-" + Path.GetFileNameWithoutExtension(gms.filePath)+"_diff_full.txt";
		File.WriteAllText(outputFile, str);
		System.Diagnostics.Process.Start(outputFile);
	}

	public void Diff(GroupMemSnapshot gms, long classSizeFilter)
    {
		if (0 > classSizeFilter)
			classSizeFilter = 0;
		
        List<KeyValuePair<string, long>> diffNativeSizeList = new List<KeyValuePair<string, long>>();
        List<KeyValuePair<string, long>> diffManageSizeList = new List<KeyValuePair<string, long>>();

        long diffTotalSize = cms.totalSize - gms.cms.totalSize;
        long diffNativeSize = cms.totalNativeObjSize - gms.cms.totalNativeObjSize;
        long diffManageSize = cms.totalManagedObjSize - gms.cms.totalManagedObjSize;

		diffClassSize (nativeSizeList, gms.nativeSizeList, classSizeFilter, ref diffNativeSizeList);
		diffClassSize (manageSizeList, gms.manageSizeList, classSizeFilter, ref diffManageSizeList);

        string str = "+++++++++++++++++++ outline diff +++++++++++++++\r\n";
        str += Path.GetFileNameWithoutExtension(filePath) + " - " + Path.GetFileNameWithoutExtension(gms.filePath) + ": \r\n";
        str += "total diff size: " +formatBytes(diffTotalSize)+"\r\n";
        str += "total diff native size: " + formatBytes(diffNativeSize) + "\r\n";
        str += "total diff manage size: " + formatBytes(diffManageSize) + "\r\n";

        str += "\r\n+++++++++++++++++++ outline snapshot +++++++++++++++\r\n";
        str += Path.GetFileNameWithoutExtension(filePath)+": "+outlineInfo()+"\r\n";
        str += Path.GetFileNameWithoutExtension(gms.filePath)+": "+gms.outlineInfo()+"\r\n";

        str += "\r\n+++++++++++++++++++++ native ++++++++++++++++\r\n";
		str += diffClassSizeStr (diffNativeSizeList, classSizeFilter, nativeDic, gms.nativeDic);

        str += "\r\n+++++++++++++++++++++ manage ++++++++++++++++\r\n";
		str += diffClassSizeStr (diffManageSizeList, classSizeFilter, manageDic, gms.manageDic);

        string outputFile = Path.GetFileNameWithoutExtension(filePath) + "-" + Path.GetFileNameWithoutExtension(gms.filePath)+"_diff.txt";
        File.WriteAllText(outputFile, str);
        System.Diagnostics.Process.Start(outputFile);
    }

    string GetGroupName(ThingInMemory thing)
    {
        if (thing is NativeUnityEngineObject)
            return (thing as NativeUnityEngineObject).className ?? "MissingName";
        if (thing is ManagedObject)
            return (thing as ManagedObject).typeDescription.name;
        return thing.GetType().Name;
    }

    string formatBytes(long size)
    {
        if (size < 0)
            return "-" + EditorUtility.FormatBytes(size * -1);
        else
            return EditorUtility.FormatBytes(size);
    }

    string outlineInfo()
    {
        return string.Format("Total memory: {0}(native: {1}, manage: {2}, gcHandle: {3}, staticFiled: {4})"
                    , EditorUtility.FormatBytes(cms.totalSize)
                    , EditorUtility.FormatBytes(cms.totalNativeObjSize)
                    , EditorUtility.FormatBytes(cms.totalManagedObjSize)
                    , EditorUtility.FormatBytes(cms.totalGcHandleSize)
                    , EditorUtility.FormatBytes(cms.totalStaticFiledSize));
    }
}
